/********************************************************************/
/*																	*/
/*  DataSocket.cpp													*/
/*																	*/
/*  Implementation of the Data Socket.								*/
/*	This class is a part of Quick 'n Easy FTP Server.				*/
/*																	*/
/*  Copyright Pablo Software Solutions 2005							*/
/*	http://www.pablosoftwaresolutions.com							*/
/*																	*/
/*  Last updated: May 28, 2005										*/
/*																	*/
/********************************************************************/


#include "stdafx.h"
#include "resource.h"
#include "DataSocket.h"
#include "ControlSocket.h"
#include "ConnectThread.h"
#include <afxpriv.h>

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

#define PACKET_SIZE 4096

/********************************************************************/
/*																	*/
/* Function name : CDataSocket::CDataSocket							*/
/* Description   : Constructor										*/
/*																	*/
/********************************************************************/
CDataSocket::CDataSocket(CControlSocket *pSocket)
{
	m_pControlSocket = pSocket;
	m_strListing = "";
	m_File.m_hFile = NULL;
	m_bConnected = FALSE;
	m_nTotalBytesSend = 0;
	m_nTotalBytesTransfered = 0;
}


/********************************************************************/
/*																	*/
/* Function name : CDataSocket::~CDataSocket						*/
/* Description   : Destructor										*/
/*																	*/
/********************************************************************/
CDataSocket::~CDataSocket()
{
	m_bConnected = FALSE;
}


#if 0
BEGIN_MESSAGE_MAP(CDataSocket, CAsyncSocket)
	//{{AFX_MSG_MAP(CDataSocket)
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()
#endif	// 0


/********************************************************************/
/*																	*/
/* Function name : OnSend											*/
/* Description   : Called by the framework to notify socket	that	*/
/*				   it can send data by calling the Send method.		*/
/*																	*/
/********************************************************************/
void CDataSocket::OnSend(int nErrorCode) 
{
	CAsyncSocket::OnSend(nErrorCode);

	switch(m_pControlSocket->m_nStatus)
	{
		case STATUS_LIST:
		{
			while (m_nTotalBytesTransfered < m_nTotalBytesSend)
			{
				DWORD dwRead;
				int dwBytes;

				CString strDataBlock;
				
				dwRead = m_strListing.GetLength();
				
				if (dwRead <= PACKET_SIZE)
				{
					strDataBlock = m_strListing;
				}
				else
				{
					strDataBlock = m_strListing.Left(PACKET_SIZE);
					dwRead = strDataBlock.GetLength();
				}
				
				if ((dwBytes = Send(strDataBlock, dwRead)) == SOCKET_ERROR)
				{
					if (GetLastError() == WSAEWOULDBLOCK) 
					{
						Sleep(0);
						return;
					}
					else
					{
						TCHAR szError[256];
						wsprintf(szError, "Server Socket failed to send: %d", GetLastError());

						// close the data connection.
						Close();

						m_nTotalBytesSend = 0;
						m_nTotalBytesTransfered = 0;

						// change status
						m_pControlSocket->m_nStatus = STATUS_IDLE;

						m_pControlSocket->SendResponse("426 Connection closed; transfer aborted.");

						// destroy this socket
						AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
					}
				}
				else
				{
					m_nTotalBytesTransfered += dwBytes;
					m_strListing = m_strListing.Mid(dwBytes);

					((CConnectThread *)AfxGetThread())->IncSentBytes(dwBytes);
				}
			}
			if (m_nTotalBytesTransfered == m_nTotalBytesSend)
			{
				// close the data connection.
				Close();

				m_nTotalBytesSend = 0;
				m_nTotalBytesTransfered = 0;

				// change status
				m_pControlSocket->m_nStatus = STATUS_IDLE;

				// tell the client the transfer is complete.
				m_pControlSocket->SendResponse("226 Transfer complete");
				// destroy this socket
				AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
			}
			break;
		}

		case STATUS_DOWNLOAD:
		{
			while (m_nTotalBytesTransfered < m_nTotalBytesSend)
			{
				// allocate space to store data
				byte data[PACKET_SIZE];
				
				m_File.Seek(m_nTotalBytesTransfered, CFile::begin);

				DWORD dwRead = m_File.Read(data, PACKET_SIZE);
    
				int dwBytes;

				if ((dwBytes = Send(data, dwRead)) == SOCKET_ERROR)
				{
					if (GetLastError() == WSAEWOULDBLOCK) 
					{
						Sleep(0);
						break;
					}
					else
					{
						TCHAR szError[256];
						wsprintf(szError, "Server Socket failed to send: %d", GetLastError());

						// close file.
						m_File.Close();
						m_File.m_hFile = NULL;

						// close the data connection.
						Close();

						m_nTotalBytesSend = 0;
						m_nTotalBytesTransfered = 0;

						// change status
						m_pControlSocket->m_nStatus = STATUS_IDLE;

						m_pControlSocket->SendResponse("426 Connection closed; transfer aborted.");

						// destroy this socket
						AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);

						// download failed
						((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_DOWNLOADFAILED);
					}
				}
				else
				{
					m_nTotalBytesTransfered += dwBytes;
					((CConnectThread *)AfxGetThread())->IncSentBytes(dwBytes);
				}
			}
			if (m_nTotalBytesTransfered == m_nTotalBytesSend)
			{
				// close file.
                m_File.Close();
				m_File.m_hFile = NULL;

                // close the data connection.
                Close();

                m_nTotalBytesSend = 0;
                m_nTotalBytesTransfered = 0;

                // change status
				m_pControlSocket->m_nStatus = STATUS_IDLE;

				// tell the client the transfer is complete.
				m_pControlSocket->SendResponse("226 Transfer complete");
				// destroy this socket
				AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
				// download successfull
				((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_DOWNLOADSUCCEEDED);
			}
			break;
		}
	}
}


/********************************************************************/
/*																	*/
/* Function name : OnAccept											*/
/* Description   : Notify this listening socket that it can accept	*/
/*				   pending connection requests.						*/
/*																	*/
/********************************************************************/
void CDataSocket::OnAccept(int nErrorCode) 
{
	// Accept the connection using a temp CSocket object.
	CAsyncSocket tmpSocket;
	Accept(tmpSocket);
	
	SOCKET socket = tmpSocket.Detach();
	Close();

	// re-use same socket for connection
	Attach(socket);

	// we're connected!
	m_bConnected = TRUE;
	
	CAsyncSocket::OnAccept(nErrorCode);
}


/********************************************************************/
/*																	*/
/* Function name : OnConnect										*/
/* Description   : Called by the framework to notify connecting		*/
/*				   socket that its connection attempt is completed	*/
/*																	*/
/********************************************************************/
void CDataSocket::OnConnect(int nErrorCode) 
{
	m_bConnected = TRUE;
	if (nErrorCode)
	{
		// change status
		m_pControlSocket->m_nStatus = STATUS_IDLE;

		m_pControlSocket->SendResponse("425 Can't open data connection.");
		// destroy this socket
		AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
	}
	CAsyncSocket::OnConnect(nErrorCode);
}


/********************************************************************/
/*																	*/
/* Function name : OnClose											*/
/* Description   : Called by the framework to notify this socket	*/
/*				   that the connected socket is closed.				*/
/*																	*/
/********************************************************************/
void CDataSocket::OnClose(int nErrorCode) 
{
	if (m_pControlSocket)
	{
		// shutdown sends
		ShutDown(1);

		if (m_pControlSocket->m_nStatus == STATUS_UPLOAD)
		{
			while(Receive() != 0)
			{
				// receive remaining data				
			}
		}
		else
		{
			m_pControlSocket->SendResponse("426 Connection closed; transfer aborted.");
			// destroy this socket
			AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
			// upload failed
			((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_UPLOADFAILED);
		}
	}
	m_pControlSocket->m_nStatus = STATUS_IDLE;	
	m_bConnected = FALSE;

	CAsyncSocket::OnClose(nErrorCode);
}


/********************************************************************/
/*																	*/
/* Function name : OnReceive										*/
/* Description   : Called by the framework to notify this socket	*/
/*				   that there is data in the buffer that can be		*/
/*				   retrieved by calling the Receive member function.*/
/*																	*/
/********************************************************************/
void CDataSocket::OnReceive(int nErrorCode) 
{
	CAsyncSocket::OnReceive(nErrorCode);
	Receive();
}


/********************************************************************/
/*																	*/
/* Function name : Receive											*/
/* Description   : Receive data from a socket						*/
/*																	*/
/********************************************************************/
int CDataSocket::Receive()
{
	int nRead = 0;
	
	if (m_pControlSocket->m_nStatus == STATUS_UPLOAD)
	{
		if (m_File.m_hFile == NULL)
			return 0;

		byte data[PACKET_SIZE];
		nRead = CAsyncSocket::Receive(data, PACKET_SIZE);

		switch(nRead)
		{
			case 0:
			{
				m_File.Close();
				m_File.m_hFile = NULL;
				Close();
				// tell the client the transfer is complete.
				m_pControlSocket->SendResponse("226 Transfer complete");
				// destroy this socket
				AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
				// upload succesfull
				((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_UPLOADSUCCEEDED);
				break;
			}
			case SOCKET_ERROR:
			{
				if (GetLastError() != WSAEWOULDBLOCK)
				{
					m_File.Close();
					m_File.m_hFile = NULL;
					Close();
					m_pControlSocket->SendResponse("426 Connection closed; transfer aborted.");
					// destroy this socket
					AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
					// upload failed
					((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_UPLOADFAILED);
				}
				break;
			}
			default:
			{
				((CConnectThread *)AfxGetThread())->IncReceivedBytes(nRead);
				TRY
				{
					m_File.Write(data, nRead);
				}
				CATCH_ALL(e)
				{
					m_File.Close();
					m_File.m_hFile = NULL;
					Close();
					m_pControlSocket->SendResponse("450 Can't access file.");
					// destroy this socket
					AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
					// upload failed
					((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_UPLOADFAILED);
					return 0;
				}
				END_CATCH_ALL;
				break;
			}
		}
	}
	return nRead;
}


/********************************************************************/
/*																	*/
/* Function name : SendListing										*/
/* Description   : Send listing to client							*/
/*																	*/
/********************************************************************/
void CDataSocket::SendListing(LPCTSTR lpszListing)
{
	m_strListing = lpszListing;
	m_nTotalBytesSend = m_strListing.GetLength();
	m_nTotalBytesTransfered = 0;
	
	// start sending...
	OnSend(0);
}


/********************************************************************/
/*																	*/
/* Function name : SendFile											*/
/* Description   : Send file to client								*/
/*																	*/
/********************************************************************/
void CDataSocket::SendFile(LPCTSTR lpszFilename)
{
	if (!PrepareSendFile(lpszFilename))
	{
		// change status
		m_pControlSocket->m_nStatus = STATUS_IDLE;

		m_pControlSocket->SendResponse("426 Connection closed; transfer aborted.");

		// destroy this socket
		AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
		return;
	}
	
	// start sending...
	OnSend(0);
}


/********************************************************************/
/*																	*/
/* Function name : RetrieveFile										*/
/* Description   : Retrieve file from client						*/
/*																	*/
/********************************************************************/
void CDataSocket::RetrieveFile(LPCTSTR lpszFilename)
{
	if (!PrepareReceiveFile(lpszFilename))
	{
		// change status
		m_pControlSocket->m_nStatus = STATUS_IDLE;

		m_pControlSocket->SendResponse("450 Can't access file.");

		// destroy this socket
		AfxGetThread()->PostThreadMessage(WM_DESTROYDATACONNECTION, 0, 0);
		// upload failed
		((CConnectThread *)AfxGetThread())->UpdateStatistic(FTPSTAT_UPLOADFAILED);
		return;
	}
}


/********************************************************************/
/*																	*/
/* Function name : PrepareReceiveFile								*/
/* Description   : Prepare socket to receive a file.				*/
/*																	*/
/********************************************************************/
BOOL CDataSocket::PrepareReceiveFile(LPCTSTR lpszFilename)
{
	// close file if it's already open
	if (m_File.m_hFile != NULL)
	{
		m_File.Close();
	}

	// open destination file
	if (!m_File.Open(lpszFilename, CFile::modeWrite | CFile::modeCreate | CFile::modeNoTruncate | CFile::shareDenyWrite))
	{
		return FALSE;
	}
	m_nTotalBytesReceive = 0;
	m_nTotalBytesTransfered = 0;
	return TRUE;
}


/********************************************************************/
/*																	*/
/* Function name : PrepareSendFile									*/
/* Description   : Prepare socket to send a file.					*/
/*																	*/
/********************************************************************/
BOOL CDataSocket::PrepareSendFile(LPCTSTR lpszFilename)
{
	// close file if it's already open
	if (m_File.m_hFile != NULL)
	{
		m_File.Close();
	}

	// open source file (bug fix by Mutex)
	if (!m_File.Open(lpszFilename, CFile::modeRead | CFile::shareDenyNone | CFile::typeBinary))
	{
		return FALSE;
	}
	m_nTotalBytesSend = m_File.GetLength();
	m_nTotalBytesTransfered = 0;
	return TRUE;
}
